#! /usr/bin/python3
"""
坐标系：屏幕向右是x+, 向上是z+, 向里是y+, 用右手展开，四指朝向屏幕，所指方向是y，大拇指方向是x，掌心向上是z。
鼠标旋转操作：比较难用，安装中键是绕原点旋转，这个原点似乎是全局的原点，不是当前屏幕原点（难用）。
"""
import logging
import os
import random
import sys
from math import cos, pi, sin
from multiprocessing import Array, Process, Queue
from threading import Thread
import time

import direct.gui.DirectGuiGlobals as DGG
import numpy as np
from direct.actor.Actor import Actor
from direct.gui.DirectGui import DirectButton, DirectEntry, DirectEntryScroll, DirectOptionMenu
from direct.gui.OnscreenText import OnscreenText
from direct.showbase.ShowBase import (Filename, LQuaternionf, LVecBase3f,
                                      NodePath, ShowBase, Task)
from IPython import embed
from panda3d.core import TextNode
from scipy.spatial.transform import Rotation

from .XmlParser import RobotTree
from .Trajectory import ComputeTrajectory
from Modules.modern_robotics import *

logger = logging.getLogger(__name__)


class Robot:
    def __init__(self, model_path, name=None, vlim=None, alim=None, interp_period=1/24):
        """Initialize a robot

        Arguments:
            model_path {str} -- robot visual model urdf path

        Keyword Arguments:
            name {str,None} -- robot name (default: {None})
            interp_period {float} -- robot control interpolation period (default: {0.01})
        """
        self.model_path = model_path
        self.name = name  # if name is None, later it will be replaced with robot name defined in urdf
        self.vlim = vlim if vlim is not None else np.r_[
            -np.ones((1, 6)), np.ones((1, 6))]
        self.alim = alim if alim is not None else np.r_[-5.0 *
                                                        np.ones((1, 6)), 5.0 * np.ones((1, 6))]
        self.interp_period = interp_period

        # [x,y,z,ox,oy,oz,ow],  position + quaternion
        self.base_transformation = [0, 0, 0, 0, 0, 0, 1]  # not copy this

        self.tree = RobotTree(model_path)

    def FK(self, joints):
        pose = FKinSpace(self.tree.m_home_pose, self.tree.m_s_list, joints)
        return pose

    def IK(self, pose):
        joints = IKinSpace(self.tree.m_s_list, self.tree.m_home_pose,
                           pose, self.GetRobotJointPosition(), 0.01, 0.001)
        return joints

    def GetName(self):
        return self.name

    def Copy(self, name):
        if name == self.name:
            logger.error("copy robot should have different name")
            return None
        robot = Robot(self.model_path, name, self.vlim,
                      self.alim, self.interp_period)
        return robot

    def SetBaseTransformation(self, pose):
        """Set robot base transformation

        Arguments:
            pose {pose} -- pose data, [x,y,z,ox,oy,oz,ow], position + quaternion

        Returns:
            bool -- result
        """
        self.base_transformation = pose
        # RobotWorld.SetRobotAttri(self, "SetBaseTransformation", pose)
        RobotWorld.robot_attribute_queue.put(
            (self.name, "SetBaseTransformation", (pose,), None))
        _, res = RobotWorld.robot_attribute_feedback_queue.get()
        return res

    def GetRobotJointPosition(self):
        return RobotWorld.robot_joints_feedback_queue.get().get(self.name, None)

    def GetRobotPose(self):
        return self.FK(self.GetRobotJointPosition())

    def MoveJ(self, joints):
        """Move robot joints

        Arguments:
            joints {list, np.array} -- target joints, dimension: dof

        Returns:
            bool -- result
        """
        traj = ComputeTrajectory(
            np.array([self.GetRobotJointPosition(), joints]), self.vlim, self.alim)
        return self.FollowTraj(traj)

    def MoveL(self, pose):
        joints, ret = self.IK(pose)
        if ret:
            return self.MoveJ(joints)
        else:
            logger.error("IK error at pose:\n{}".format(pose))
            return False

    def FollowTraj(self, traj):
        """Follow a toppra trajecotry

        Arguments:
            traj {class} -- joint trajectory, should have property "duration" and "eval" method

        Returns:
            bool -- result
        """
        t0 = time.time()
        t = 0
        while t < traj.duration:
            if not self._SetRobotJointPosition(traj.eval(t)):
                logger.error("Follow trajectory error at {}".format(t))
                return False
            t = time.time() - t0
            time.sleep(self.interp_period)
        return True

    def _SetRobotJointPosition(self, joints):
        joints = {self.name: joints}
        RobotWorld.robot_joints_feed_queue.put(joints)
        _, res = RobotWorld.robot_joints_feed_reply_queue.get()
        return res

    def __eq__(self, robot):
        return self.name == robot.name

    def __repr__(self):
        return "<Robot at {:#x}, name: {}>".format(id(self), self.name)


class ZenWorld(ShowBase):
    def __init__(self):
        ShowBase.__init__(self)
        self.set_background_color(0.9, 0.9, 0.9, 0)


class RobotWorld(ZenWorld):
    class _Robot:
        def __init__(self, model_path):
            self.initial_joint_state = [0, 0, 0, 0, 0, 0]
            self.joint_feedback = self.initial_joint_state
            self.joint_feed = self.initial_joint_state
            self.joints = None
            self.link_nodes = []
            if model_path:
                self.robot_tree = RobotTree(model_path)
                self.dof = self.robot_tree.GetDoF()
                self.name = self.robot_tree.GetName()
            else:
                self.robot_tree = None

        def SetBaseTransformation(self, pose):
            pos, quat = pose[:3], pose[3:]
            hpr = Rotation.from_quat(quat).as_euler("zxy")
            self.link_nodes[1].set_pos(*pos)
            self.link_nodes[1].set_hpr(*hpr)
            return True

        def __repr__(self):
            return "<_Robot at {:#x}, name: {}>".format(id(self), self.name)

    robots = {}
    add_robot = None
    remove_robot = None
    load_robot_path_queue = Queue()
    load_robot_feedback_queue = Queue()
    remove_robot_name_queue = Queue()
    remove_robot_feedback_queue = Queue()
    robot_joints_feed_queue = Queue()
    robot_joints_feed_reply_queue = Queue()
    robot_joints_feedback_queue = Queue()
    robot_attribute_queue = Queue()
    robot_attribute_feedback_queue = Queue()
    world_attribute_queue = Queue()
    world_attribute_feedback_queue = Queue()

    def __init__(self):
        ZenWorld.__init__(self)
        # self.m_models = [self.render]

        # load robot from urdf
        # robot_tree = RobotTree(model_path)
        # self.robot_tree = robot_tree
        # rgba = robot_tree.GetRGBA()
        # self.joints = robot_tree.GetJoints()
        # self.links = robot_tree.GetLinks()
        # for node in robot_tree.GetChain():
        #     link = self.links.get(node, False)
        #     if not link or link.get("visual_mesh", "") == "":
        #         continue
        #     model = self.loader.loadModel(
        #         Filename.from_os_specific(link["visual_mesh"]))
        #     model.reparent_to(self.m_models[-1])
        #     model.set_color(*rgba)
        #     self.m_models.append(model)
        #     if link["parent"]:
        #         parent_joint_pos = self.joints[link["parent"]]["origin_xyz"]
        #         parent_joint_rpy = self.joints[link["parent"]]["origin_rpy"]
        #         model.set_pos(*parent_joint_pos)
        #         model.set_hpr(*np.rad2deg(parent_joint_rpy))
        self.joints_move_direction = [0, 0, 0, 0, 0, 0]

        # camera control ################3

        self.disable_mouse()
        self.oobe(cam=self.camera)
        self.camera.set_pos(0, -4, 4)
        self.camera.lookAt(self.render)

        self.key_map = {}
        self.accept("a", self._SetKey, ["cam_xminus", True])
        self.accept("a-up", self._SetKey, ["cam_xminus", False])
        self.accept("d", self._SetKey, ["cam_xplus", True])
        self.accept("d-up", self._SetKey, ["cam_xplus", False])
        self.accept("w", self._SetKey, ["cam_zplus", True])
        self.accept("w-up", self._SetKey, ["cam_zplus", False])
        self.accept("s", self._SetKey, ["cam_zminus", True])
        self.accept("s-up", self._SetKey, ["cam_zminus", False])
        self.accept("q", self._SetKey, ["cam_yplus", True])
        self.accept("wheel_up", self._SetKey, ["cam_yplus", True])
        # self.accept("wheel_up", lambda: print("wheel up"))
        self.accept("q-up", self._SetKey, ["cam_yplus", False])
        self.accept("wheel_down", self._SetKey, ["cam_yminus", True])
        # self.accept("wheel_down", lambda: print("wheel down"))
        self.accept("e", self._SetKey, ["cam_yminus", True])
        self.accept("e-up", self._SetKey, ["cam_yminus", False])
        # self.accept("arrow_up", self._SetKey, ["cam_zplus", True])
        # self.accept("arrow_up-up", self._SetKey, ["cam_zplus", False])
        # self.accept("arrow_down", self._SetKey, ["cam_zminus", True])
        # self.accept("arrow_down-up", self._SetKey, ["cam_zminus", False])
        # self.accept("i", self._SetKey, ["cam_pplus", True])
        # self.accept("i-up", self._SetKey, ["cam_pplus", False])
        # self.accept("k", self._SetKey, ["cam_pminus", True])
        # self.accept("k-up", self._SetKey, ["cam_pminus", False])
        # self.accept("arrow_left", self._SetKey, ["cam_xminus", True])
        # self.accept("arrow_left-up", self._SetKey, ["cam_xminus", False])
        # self.accept("arrow_right", self._SetKey, ["cam_xplus", True])
        # self.accept("arrow_right-up", self._SetKey, ["cam_xplus", False])
        self.accept("space", self._SetKey, ["cam_origin", True])
        self.accept("space-up", self._SetKey, ["cam_origin", False])
        self.accept("space-up", self._SetKey, ["cam_xplus", False])
        # self.accept("escape", sys.exit, [0])
        self.accept("p", self._ExecCmdNewThrd)
        self.accept("f1", self._ToggleJointControlPanel)

        self._AddInstructions(0.05, "[f1]: show/toggle joint control panel")
        self._AddInstructions(0.1, "[space]: look at origin")
        self._AddInstructions(0.15, "[q/e]: zoom in/out")
        self._AddInstructions(0.2, "[w/s/a/d]: move camera")
        self._AddInstructions(
            0.25, "[mouse]: press middle button to rotate camera")
        self.joint_panel_show = False

        # # Add the spinCameraTask procedure to the task manager.
        # self.taskMgr.add(self.spinCameraTask, "SpinCameraTask")
        self.taskMgr.add(self._UpdateCameraTask, "UpdateCameraTask")
        self.taskMgr.add(self._UpdateJointJogTask, "UpdateJointJogTask")
        self.taskMgr.add(self._UpdateRobotTask, "UpdateRobotTask")
        self.move_joints_panel_frame = []
        self.robot_menu = None
        # entry = DirectEntry(parent=self.a2dBottomLeft, text="",
        #                     pos=LVecBase3f(0.1, 0.2), scale=0.5, text_scale=0.5)
        # self.editor = DirectEntryScroll(entry=entry, parent=self.a2dBottomLeft)

    def _AddInstructions(self, pos, msg):
        return OnscreenText(text=msg, style=1, fg=(0, 0, 1, 1), scale=.04,
                            shadow=(0, 0, 0, 0), parent=base.a2dBottomLeft,
                            pos=(0.08, pos), align=TextNode.ABoxedLeft)

    def _ToggleJointControlPanel(self):
        if self.joint_panel_show:
            self._DestroyJointControlPanel()
            self.joint_panel_show = False
        else:
            self._ShowJointControlPanel()
            self.joint_panel_show = True

    def _DestroyJointControlPanel(self):
        for button in self.move_joints_panel_frame:
            button.hide()
        self.robot_menu.hide()

    def _ShowJointControlPanel(self):
        if len(self.move_joints_panel_frame) == 0:
            for i in range(6):
                self.move_joints_panel_frame.append(DirectButton(text=("J{}+".format(i+1), "Moving", "J{}+".format(i+1), "disabled"),
                                                                 scale=.08, command=lambda x=i: self._SetJointMoveDirection(x, 1),
                                                                 pos=LVecBase3f(0.1, -0.1 - i * 0.11), relief=DGG.GROOVE,
                                                                 parent=self.a2dTopLeft, frameSize=(-1.5, 1.5, -0.4, 0.9)))
            for i in range(6):
                self.move_joints_panel_frame.append(DirectButton(text=("J{}-".format(i+1), "Moving", "J{}-".format(i+1), "disabled"),
                                                                 scale=.08, command=lambda x=i: self._SetJointMoveDirection(x, -1),
                                                                 pos=LVecBase3f(0.32, -0.1 - i * 0.11), relief=DGG.GROOVE,
                                                                 parent=self.a2dTopLeft, frameSize=(-1.5, 1.5, -0.4, 0.9)))
            for i in range(6):
                self.move_joints_panel_frame.append(DirectButton(text=("Stop J{}".format(i+1), "Stop", "Stop J{}".format(i+1), "disabled"),
                                                                 scale=.08, command=lambda x=i: self._SetJointMoveDirection(x, 0),
                                                                 pos=LVecBase3f(0.6, -0.1 - i * 0.11), relief=DGG.GROOVE,
                                                                 parent=self.a2dTopLeft, frameSize=(-2.0, 2.0, -0.4, 0.9)))
        if self.robots:
            self.robot_menu = DirectOptionMenu(items=list(self.robots.keys()), scale=0.1, pos=LVecBase3f(0.001, 0.2),
                                               text_scale=0.8, item_text_scale=1.0, text_align=TextNode.ABoxedLeft,
                                               relief=DGG.GROOVE, parent=self.a2dLeftCenter, frameSize=(-0.0, 7.0, -0.4, 0.9))
        for button in self.move_joints_panel_frame:
            button.show()

    def _UpdateCameraTask(self, task):
        # time since last frame
        dt = globalClock.getDt()
        camera_current_pos = self.camera.get_pos()
        step = (camera_current_pos.length() + 0.1) * 0.1
        zoom_ratio = 20.0
        if self.key_map.get("cam_xplus", False):
            self.camera.setX(self.camera, step*dt)
        elif self.key_map.get("cam_xminus", False):
            self.camera.setX(self.camera, -step*dt)
        elif self.key_map.get("cam_yplus", False):
            self.camera.setY(self.camera, step*dt*zoom_ratio)
            self.key_map["cam_yplus"] = False
        elif self.key_map.get("cam_yminus", False):
            self.camera.setY(self.camera, -step*dt*zoom_ratio)
            self.key_map["cam_yminus"] = False
        elif self.key_map.get("cam_zplus", False):
            self.camera.setZ(self.camera, step*dt)
        elif self.key_map.get("cam_zminus", False):
            self.camera.setZ(self.camera, -step*dt)
        elif self.key_map.get("cam_hplus", False):
            self.camera.setH(self.camera, step*dt)
        elif self.key_map.get("cam_hminus", False):
            self.camera.setH(self.camera, -step*dt)
        elif self.key_map.get("cam_pplus", False):
            self.camera.setP(self.camera, step*dt)
        elif self.key_map.get("cam_pminus", False):
            self.camera.setP(self.camera, -step*dt)
        elif self.key_map.get("cam_rplus", False):
            self.camera.setR(self.camera, step*dt)
        elif self.key_map.get("cam_rminus", False):
            self.camera.setR(self.camera, -step*dt)
        elif self.key_map.get("cam_origin", False):
            # self.camera.set_pos(0, -4, 0)
            # self.camera.set_hpr(0, 0, 0)
            self.camera.lookAt(self.render)
        return Task.cont

    def _UpdateJointJogTask(self, task):
        if hasattr(self, "robot_menu") and self.robot_menu is not None:
            robot = self.robots[self.robot_menu.get()]
            step = 0.4
            for i, direction in enumerate(self.joints_move_direction):
                if direction > 0:
                    self._MoveSingleJoint(robot, step, i)
                elif direction < 0:
                    self._MoveSingleJoint(robot, -step, i)
        return Task.cont

    def _ExecCmdNewThrd(self):
        Thread(target=self._ExecCmd).start()

    def _SetKey(self, key, val):
        self.key_map[key] = val

    def _SetJointMoveDirection(self, joint_num, direction):
        self.joints_move_direction[joint_num] = direction

    def _ExecCmd(self):
        while True:
            s = input("input cmd:")
            if s == "exit":
                print("bye!")
                break
            try:
                print(eval(s))
            except Exception as e:
                print(e)

    def _SpinCameraTask(self, task: Task):
        angleDegrees = task.time * 6.0
        angleRadians = angleDegrees * (pi / 180.0)
        distance = 10
        x = distance * sin(angleRadians)
        y = - distance * cos(angleRadians)
        z = distance
        self.camera.setPos(x, y, z)
        self.camera.setHpr(angleDegrees, -45, 0)
        return Task.cont

    ################# Robot control api ####################

    def _LoadRobot(self):
        if self.add_robot:
            robot = self.add_robot
            self.add_robot = None
            if robot.name in self.robots:
                logger.error(
                    "robot {} is already in RobotWorld, use another name and try again".format(robot.name))
                self.load_robot_feedback_queue.put((robot.name, False))
                return
            self.robots[robot.name] = robot
            robot.link_nodes = [self.render]
            # load robot from urdf
            robot_tree = robot.robot_tree
            rgba = robot_tree.GetRGBA()
            robot.joints = robot_tree.GetJoints()
            robot.links = robot_tree.GetLinks()
            for node in robot_tree.GetChain():
                link = robot.links.get(node, False)
                if not link or link.get("visual_mesh", "") == "":
                    continue
                node = self.loader.loadModel(
                    Filename.from_os_specific(link["visual_mesh"]))
                node.reparent_to(robot.link_nodes[-1])
                node.set_color(*rgba)
                robot.link_nodes.append(node)
                if link["parent"]:
                    parent_joint_pos = robot.joints[link["parent"]
                                                    ]["origin_xyz"]
                    parent_joint_rpy = robot.joints[link["parent"]
                                                    ]["origin_rpy"]
                    node.set_pos(*parent_joint_pos)
                    node.set_hpr(*np.rad2deg(parent_joint_rpy))
            robot.joints_move_direction = [0, 0, 0, 0, 0, 0]
            self.load_robot_feedback_queue.put((robot.name, True))

    def _RemoveRobot(self):
        if self.remove_robot:
            robot = self.remove_robot
            self.remove_robot = None
            if robot.name not in self.robots:
                logger.error("robot {} not in RobotWorld".format(robot.name))
                self.remove_robot_name_queue.put((robot.name, False))
                return
            self.robots.pop(robot.name)
            # the first node is self.render, don't remove it
            node_number = len(robot.link_nodes)
            while node_number > 1:
                robot.link_nodes.pop().remove_node()
                node_number -= 1
            self.remove_robot_feedback_queue.put((robot.name, True))

    def _MoveSingleJoint(self, robot, degree=1.0, joint_num=0, absolute=False):
        """This is most base api to move robot, all other move functions will finally call this one.

        Keyword Arguments:
            degree {float} -- how many degrees to move (default: {1})
            joint_num {int} -- the jonit index, from 0 to dof-1 (default: {0})
            absolute {bool} -- move the joint relative current position or an absolute position (default: {False})
        """
        J = robot.joints[robot.robot_tree.joint_names[joint_num]]
        x, y, z = J.get("axis", np.zeros(3))
        # the euler sequence is quite odd
        joints_inc = LVecBase3f(z, x, y) * degree
        # skip self.render node and base link node
        node_J = robot.link_nodes[2+joint_num]
        if absolute:
            node_J.setHpr(joints_inc)
            robot.joint_feedback[joint_num] = np.deg2rad(degree)
        else:
            node_J.setHpr(node_J, joints_inc)
            robot.joint_feedback[joint_num] += np.deg2rad(degree)

    def _SetJointPositions(self, robot, joints):
        """ Set robot joint to given joint value immediately

        Arguments:
            joints {list,np.array} -- value of each joint in rad
        """
        for i, joint in enumerate(joints):
            self._MoveSingleJoint(robot, np.rad2deg(joint), i, absolute=True)
        self.robot_joints_feed_reply_queue.put((robot.name, True))

    def _UpdateJointPosition(self):
        if not self.robot_joints_feed_queue.empty():
            for name, joints in self.robot_joints_feed_queue.get().items():
                robot = self.robots.get(name, None)
                if robot:
                    self._SetJointPositions(robot, joints)

    def _UpdateRobotTask(self, task):
        if not self.load_robot_path_queue.empty():
            robot = self.load_robot_path_queue.get()
            self.add_robot = RobotWorld._Robot(robot.model_path)
            if robot.name:
                self.add_robot.name = robot.name
            else:
                robot.name = self.add_robot.name
        if not self.remove_robot_name_queue.empty():
            self.remove_robot = self.robots.get(
                self.remove_robot_name_queue.get(), None)
        self._RemoveRobot()
        self._LoadRobot()
        self._UpdateJointPosition()
        self._RobotAttribute()
        joint_feedback = {}
        for name, robot in self.robots.items():
            joint_feedback[name] = robot.joint_feedback
        if not self.robot_joints_feedback_queue.empty():
            self.robot_joints_feedback_queue.get_nowait()
        self.robot_joints_feedback_queue.put(joint_feedback)
        return Task.cont

    def _SetRobotAttr(self):
        """
        - SetRobotJointPosition  OK
        - SetRobotName
        - SetRobotBaseTransformation
        """
        pass

    def _GetRobots(self):
        return list(self.robots.keys())

    def _RobotAttribute(self):
        if not self.robot_attribute_queue.empty():
            owner_name, attribute, args, kwargs = self.robot_attribute_queue.get()
            if owner_name != "world":
                owner = self.robots.get(owner_name)
            else:
                owner = self
            if owner:
                logger.debug("args: {}, kwargs: {}".format(args, kwargs))
                if hasattr(owner, attribute):
                    attri = getattr(owner, attribute)
                    val = False
                    if callable(attri):
                        if args and kwargs:
                            val = attri(*args, **kwargs)
                        elif args:
                            val = attri(*args)
                        elif kwargs:
                            val = attri(**kwargs)
                        else:
                            val = attri()
                        self.robot_attribute_feedback_queue.put(
                            (owner_name, val))
                    else:
                        attri = args
                else:
                    self.robot_attribute_feedback_queue.put(
                        (owner_name, False))


class RobotLoader:
    robots = {}
    @classmethod
    def LoadRobot(cls, robot):
        RobotWorld.load_robot_path_queue.put(robot)
        name, result = RobotWorld.load_robot_feedback_queue.get()  # feedback name
        robot.name = name
        if result:
            cls.robots[robot.name] = robot
        return result

    @classmethod
    def RemoveRobot(cls, robot):
        RobotWorld.remove_robot_name_queue.put(robot.name)
        _, result = RobotWorld.remove_robot_feedback_queue.get()  # True or False
        if result:
            cls.robots.pop(robot.name)
        return result

    @classmethod
    def GetRobots(cls):
        return cls.robots


def StartRobotWorld():
    Process(target=lambda: RobotWorld().run(), daemon=True).start()
    return RobotLoader


if __name__ == "__main__":
    CWD = os.path.dirname(__file__)
    upper_dir = os.path.dirname(CWD)
    model_path = os.path.join(upper_dir, "Models/GP12/GP12.urdf")
    model_path = os.path.join(upper_dir, "Models/Motomini/Motomini.urdf")
    app = RobotWorld()
    app.run()
